from neo4j import GraphDatabase
import requests
import json
import time
import logging
from dataclasses import dataclass
from typing import List, Dict, Optional
import re


# === 配置部分 ===
@dataclass
class Config:
    NEO4J_URI = "bolt://localhost:7687"
    NEO4J_AUTH = ("neo4j", "12345678")
    OLLAMA_ENDPOINT = "http://localhost:11434/api/generate"
    CHAT_MODEL = "gemma3:4b"
    MIN_SIMILARITY = 0.7  # 最小相似度阈值


# === 知识图谱增强模块 ===
class KnowledgeEnhancer:
    def __init__(self, uri: str, auth: tuple):
        self.driver = GraphDatabase.driver(uri, auth=auth)

    def check_database_content(self, query: str) -> Optional[Dict]:
        """检查输入是否可能是数据库内容"""
        # 先尝试精确匹配
        exact_match = self._query_exact_match(query)
        if exact_match:
            return exact_match

        # 模糊匹配
        return self._query_fuzzy_match(query)

    def _query_exact_match(self, name: str) -> Optional[Dict]:
        """精确查询"""
        cypher = """
        MATCH (m)
        WHERE m.name = $name AND ANY(label IN labels(m) WHERE label IN ['方剂', '药材'])
        RETURN 
            m.name as name,
            labels(m) as type,
            CASE 
                WHEN '方剂' IN labels(m) THEN m.consist
                ELSE null
            END as details
        LIMIT 1
        """
        with self.driver.session() as session:
            result = session.run(cypher, name=name)
            record = result.single()
            return dict(record) if record else None

    def _query_fuzzy_match(self, text: str) -> Optional[Dict]:
        """模糊查询"""
        cypher = """
        MATCH (m)
        WHERE ANY(label IN labels(m) WHERE label IN ['方剂', '药材'])
        AND m.name CONTAINS $text
        RETURN 
            m.name as name,
            labels(m) as type,
            CASE 
                WHEN '方剂' IN labels(m) THEN m.consist
                ELSE null
            END as details
        ORDER BY m.name
        LIMIT 1
        """
        with self.driver.session() as session:
            result = session.run(cypher, text=text)
            record = result.single()
            return dict(record) if record else None

    def get_related_info(self, name: str) -> Dict:
        """获取完整关联信息"""
        cypher = """
        MATCH (m {name: $name})-[r]-(related)
        WHERE ANY(label IN labels(m) WHERE label IN ['方剂', '药材'])
        RETURN 
            m.name as subject,
            labels(m) as subject_type,
            type(r) as relation_type,
            labels(related) as related_type,
            related.name as related_name
        ORDER BY relation_type
        """
        with self.driver.session() as session:
            result = session.run(cypher, name=name)
            return {
                "subject": name,
                "relations": [dict(record) for record in result]
            }


# === AI对话引擎 ===
class AIChatEngine:
    def __init__(self, config: Config):
        self.config = config
        self.session = requests.Session()
        self._setup_session()

    def _setup_session(self):
        retry_strategy = requests.adapters.Retry(
            total=3, backoff_factor=0.5,
            status_forcelist=[408, 429, 500, 502, 503, 504]
        )
        adapter = requests.adapters.HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

    def stream_response(self, prompt: str):
        """流式获取AI响应"""
        try:
            response = self.session.post(
                self.config.OLLAMA_ENDPOINT,
                json={
                    "model": self.config.CHAT_MODEL,
                    "prompt": prompt,
                    "stream": True
                },
                stream=True,
                timeout=300
            )
            response.raise_for_status()

            for chunk in response.iter_lines():
                if chunk:
                    try:
                        data = json.loads(chunk.decode())
                        if not data.get("done", False) and "response" in data:
                            yield data["response"]
                    except json.JSONDecodeError:
                        continue

        except requests.exceptions.RequestException as e:
            yield f"[系统错误] AI服务不可用: {str(e)}"


# === 主控系统 ===
class MedicineAssistant:
    def __init__(self):
        self.config = Config()
        self.knowledge = KnowledgeEnhancer(self.config.NEO4J_URI, self.config.NEO4J_AUTH)
        self.ai = AIChatEngine(self.config)
        self._setup_logging()

    def _setup_logging(self):
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[logging.StreamHandler()]
        )

    def _format_context(self, data: Dict) -> str:
        """格式化知识图谱数据为上下文"""
        if not data:
            return "无相关数据库信息"

        relations = data.get("relations", [])
        if not relations:
            return f"数据库中有 {data['subject']} 但无关联信息"

        grouped = {}
        for rel in relations:
            key = rel['relation_type']
            if key not in grouped:
                grouped[key] = []
            grouped[key].append(
                f"{rel['related_name']} ({'|'.join(rel['related_type'])})"
            )

        lines = [f"【{data['subject']}】相关数据："]
        for rel_type, items in grouped.items():
            lines.append(f"{rel_type}:")
            lines.extend([f"  - {item}" for item in items])

        return "\n".join(lines)

    def process_query(self, query: str):
        """智能处理用户查询"""
        # 检查是否为数据库内容
        db_match = self.knowledge.check_database_content(query)

        if db_match:
            print(f"\n[系统] 检测到数据库内容: {db_match['name']} ({'|'.join(db_match['type'])})")

            # 获取完整关联信息
            full_data = self.knowledge.get_related_info(db_match['name'])
            context = self._format_context(full_data)

            print("\n[知识图谱数据]")
            print(context)

            # 生成AI分析
            print("\n[AI分析]")
            self._generate_response(
                f"请综合分析以下中药方剂信息: {db_match['name']}",
                context
            )
        else:
            # 普通问答模式
            print("\n[AI回答]")
            self._generate_response(query)

    def _generate_response(self, query: str, context: str = ""):
        """生成AI响应"""
        prompt = f"""
        你是一名专业中医助手，请根据以下信息回答问题。
        注意：回答需专业准确，使用中文，分点说明。

        {f"[知识库信息]\n{context}" if context else "[无相关数据库信息]"}

        [用户问题]
        {query}
        """

        print("AI: ", end="", flush=True)
        for chunk in self.ai.stream_response(prompt):
            print(chunk, end="", flush=True)
            time.sleep(0.03)
        print()

    def run(self):
        """运行交互系统"""
        print("=== 中医智能辅助系统 ===")
        print("可输入中药/方剂名称获取专业分析，或直接提问\n")

        while True:
            try:
                user_input = input("\n请输入> ").strip()
                if user_input.lower() in ('exit', 'quit'):
                    break

                self.process_query(user_input)

            except KeyboardInterrupt:
                print("\n正在退出系统...")
                break
            except Exception as e:
                print(f"\n[错误] 系统异常: {str(e)}")


if __name__ == "__main__":
    assistant = MedicineAssistant()
    assistant.run()